"""
优化的边特征聚合操作

使用PyTorch原生操作实现高效的边特征聚合，避免Python循环
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
import sys
import os

# 添加路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(os.path.dirname(current_dir))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

from ops.src.graph.graph import is_dgl_graph


def copy_e_sum_optimized(graph, edge_features):
    """
    优化的边特征聚合到目标节点操作
    
    参数:
        graph: Graph对象或DGL图
        edge_features: 边特征张量
    
    返回:
        聚合后的节点特征
    """
    print("===================copy_e_sum_optimized===================")
    # 检查图是否为DGL图
    if is_dgl_graph(graph):
        # DGL图 - 使用原生DGL操作
        import dgl.function as fn
        graph = graph.local_var()
        graph.edata['e'] = edge_features
        graph.update_all(fn.copy_e('e', 'm'), fn.sum('m', 'out'))
        return graph.ndata['out']
    else:
        # PyG风格的Graph - 使用torch_scatter高效实现
        dst = graph.edge_index[1]
        
        # 确保num_nodes是整数
        if hasattr(graph, 'num_nodes'):
            if callable(graph.num_nodes):
                num_nodes = graph.num_nodes()
            else:
                num_nodes = graph.num_nodes
        else:
            # 如果没有num_nodes属性，则从边索引中推断
            num_nodes = int(dst.max().item() + 1)
        
        # 使用torch_scatter进行高效聚合
        try:
            from torch_scatter import scatter_add
            # scatter_add is much faster than manual loops
            if edge_features.dim() == 1:
                # 1D features
                result = scatter_add(edge_features, dst, dim=0, dim_size=num_nodes)
            else:
                # Multi-dimensional features
                result = scatter_add(edge_features, dst, dim=0, dim_size=num_nodes)
            return result
        except ImportError:
            # Fallback to torch native implementation if torch_scatter not available
            return _copy_e_sum_torch_native(dst, edge_features, num_nodes)


def _copy_e_sum_torch_native(dst, edge_features, num_nodes):
    """
    使用PyTorch原生操作的高效实现（无需torch_scatter）
    """
    dtype = edge_features.dtype
    device = edge_features.device
    
    # 创建输出张量
    if edge_features.dim() == 1:
        result = torch.zeros(num_nodes, dtype=dtype, device=device)
    else:
        output_shape = [num_nodes] + list(edge_features.shape[1:])
        result = torch.zeros(output_shape, dtype=dtype, device=device)
    
    # 使用index_add_进行高效聚合（比循环快得多）
    if edge_features.dim() == 1:
        result.index_add_(0, dst, edge_features)
    else:
        result.index_add_(0, dst, edge_features)
    
    return result 